import numpy as np
import torch as th
from torch import nn
from torchvision.models.resnet import resnet18, ResNet18_Weights
from torchvision.transforms import v2

conv1 = nn.Conv2d(3, 3, 3, stride=1, padding=1)

resnet = resnet18(weights=ResNet18_Weights.DEFAULT)
dual_input = th.randn(2, 3, 224, 224, dtype=th.float32)
_ = conv1(dual_input)

class Preprocess(nn.Module):

    def __init__(self, bulk: nn.Module):
        super().__init__()
        self.bulk = bulk
    def forward(self, x):
        mean = th.tensor([0.485, 0.456, 0.406], dtype=th.float32, device=x.device)
        std = th.tensor([0.229, 0.224, 0.225], dtype=th.float32, device=x.device)
        norm = (x - mean.expand(x.size())) / std.expand(x.size())
        return self.bulk(norm)

normalize = v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

to_dump = nn.Sequential(normalize, resnet, nn.Softmax(dim=1))
to_dump.eval()
with th.no_grad():
    output = to_dump(dual_input)
    output = to_dump(dual_input)
    print('output.shape:', output.shape)
    print('th.sum(output):', th.sum(output, dim=1))
    print('th.sum(output):', th.sum(output, dim=1))

#to_dump = resnet
#to_dump = Preprocess(resnet)

BATCH_SIZE = 1
dummy_input = th.randn(BATCH_SIZE, 3, 224, 224)

th.onnx.export(to_dump, dummy_input, 'resnet18_softmax.onnx', verbose=True)



